json = require "cjson"
ngx.update_time()
local server = require "resty.websocket.server"
local wb, err = server:new{
    timeout = 50000,  -- in milliseconds
    max_payload_len = 6553500,
}
if not wb then
    ngx.log(ngx.ERR, "failed to new websocket: ", err)
    return ngx.exit(444)
end
function read_network(types)
    file_rx  = io. open ( "/sys/class/net/eth0/statistics/rx_bytes" , "r" )
    file_tx  = io. open ( "/sys/class/net/eth0/statistics/tx_bytes" , "r" )
    local  tx1 =  file_tx : read ( "*l" )
    local  rx1 =  file_rx : read ( "*l" )
    ngx.sleep (1)
    file_rx :seek( "set" )
    file_tx :seek( "set" )
    local  tx2 =  file_tx : read ( "*l" )
    local  rx2 =  file_rx : read ( "*l" )
    local  tx = (tx2-tx1)*8
    local  rx = (rx2-rx1)*8
    local total = tx+rx
    mb = {time=ngx.now()}
    if (types == 1)
    then
        if (rx<10000)
        then
            mb["rx"]=rx
            mb["rx_mode"]="bps"
        elseif(rx<10000000)
        then
            rx = rx/1024
            mb["rx"]=rx
            mb["rx_mode"]="Kbps"
        elseif(rx<10000000000)
        then
            rx = rx/1024/1024
            mb["rx"]=rx
            mb["rx_mode"]="Mbps"
        elseif(rx<10000000000000)
        then
            rx = rx/1024/1024/1024
            mb["rx"]=rx
            mb["rx_mode"]="Gbps"
        end
        
        if (tx<10000)
        then
            mb["tx"]=tx
            mb["tx_mode"]="bps"
        elseif(tx<10000000)
        then
            tx = tx/1024
            mb["tx"]=tx
            mb["tx_mode"]="Kbps"
        elseif(tx<10000000000)
        then
            tx = tx/1024/1024
            mb["tx"]=tx
            mb["tx_mode"]="Mbps"
        elseif(tx<10000000000000)
        then
            tx = tx/1024/1024/1024
            mb["tx"]=tx
            mb["tx_mode"]="Gbps"
            
        end
        
        if (total<10000)
        then
            mb["total"] = total
            mb["total_mode"]="bps"
        elseif(total<10000000)
        then
            total = total/1024
            mb["total"] = total
            mb["total_mode"]="Kbps"
        elseif(total<10000000000)
        then
            total = total/1024/1024
            mb["total"] = total
            mb["total_mode"]="Mbps"
        elseif(total<10000000000000)
        then
            total = total/1024/1024/1024
            mb["total"] = total
            mb["total_mode"]="Gbps"
            
        end
        
    elseif(types == 2)
    then
        mb["tx"]=tx/1024/1024
        mb["tx_mode"]="Mbps"
        mb["rx"]=rx/1024/1024
        mb["rx_mode"]="Mbps"
        mb["total"] = total/1024/1024
        mb["total_mode"]="Mbps"
    else
        mb["tx"]=tx
        mb["tx_mode"]="bps"
        mb["rx"]=rx
        mb["rx_mode"]="bps"
    end
    file_rx :close()
    file_tx :close()
    return mb
end
local types = tonumber(ngx.var.arg_types) or 0
local pod = ngx.var.arg_pod or "tx"

bytes, err = wb:send_text(json.encode({Msg="Connect OK,Service will traceback the traffic data soon",code = "OK"}))
if not bytes then
        ngx.log(ngx.ERR, "failed to send a text frame: ", err)
end
ngx.sleep(1)
while (1)
do
    bytes, err = wb:send_text(json.encode(read_network(types,pod)))
    if not bytes then
        ngx.log(ngx.ERR, "failed to send a text frame: ", err)
    end
end